{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9da9a29c",
   "metadata": {},
   "source": [
    "# Dynamic Workflows\n",
    "\n",
    "To dynamically create an electron based on the output of another electron, use a sublattice.\n",
    "\n",
    "## Context\n",
    "\n",
    "A sublattice is a `lattice` decorated with an `electron`. All the restrictions of a lattice apply to a sublattice. Most importantly, computations, especially result-dependent ones, should be carried out inside electrons to reduce the danger of an error when constructing the transport graph. \n",
    "\n",
    "Since it is also an electron, a sublattice is executed as a part of the workflow. Because of this dual identity, dynamic code can be run inside a sublattice which would otherwise be impossible in the Covalent framework. For example, sublattices provide a way to handle result-dependent loops and if/else statements.\n",
    "\n",
    "## Best Practice\n",
    "\n",
    "A sublattice enables you to compose and encapsulate arbitrarily complex code. Use a sublattice (a `lattice` decorated with an `electron`) to encapsulate dynamic code that would otherwise be difficult or impossible to execute correctly in the Covalent paradigm. Use the same best practices in building a sublattice that you would for any other lattice: [Do all calculations in electrons](./post_process.ipynb); [don't create unnecessary executors](./executor_assignment.ipynb); and so on.\n",
    "\n",
    "You can nest sublattices to any level.\n",
    "\n",
    "\n",
    "## Example\n",
    "\n",
    "Contrast the two examples below.\n",
    "\n",
    "### Example 1: Incorrect\n",
    "\n",
    "The following example contains code in the `workflow_1` lattice that is not inside electrons: in this case, an if/else decision and a `for` loop. The workflow fails when Covalent tries to build the transport graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7aaeba99",
   "metadata": {},
   "outputs": [],
   "source": [
    "import covalent as ct\n",
    "\n",
    "# Technique 1: Incorrect\n",
    "\n",
    "@ct.electron\n",
    "def task_1(x):\n",
    "    return x * 3\n",
    "\n",
    "@ct.electron\n",
    "def task_2(x):\n",
    "    return x ** 2\n",
    "\n",
    "@ct.lattice\n",
    "def workflow_1(a):\n",
    "    \n",
    "    res = task_1(a)\n",
    "    \n",
    "    # An if/else decision and a result-dependent loop with no enclosing electron\n",
    "    if res < 10: \n",
    "        final_res = []\n",
    "        for _ in range(res):\n",
    "            final_res.append(task_2(res))\n",
    "    else:\n",
    "        final_res = res\n",
    "    \n",
    "    return final_res\n",
    "\n",
    "# Uncomment these three lines to see the workflow fail:\n",
    "# id = ct.dispatch(workflow_1)(2)\n",
    "# res = ct.get_result(id, wait=True)\n",
    "# print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4db4a2b",
   "metadata": {},
   "source": [
    "### Example 2: Correct\n",
    "\n",
    "The following code corrects the previous example by enclosing the if/else decision and the `for` loop in the `sub_workflow` sublattice. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "474717f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Lattice Result\n",
      "==============\n",
      "status: COMPLETED\n",
      "result: [0, 1, 4, 9, 16, 25]\n",
      "input args: ['2']\n",
      "input kwargs: {}\n",
      "error: None\n",
      "\n",
      "start_time: 2023-03-16 19:16:18.195332\n",
      "end_time: 2023-03-16 19:16:18.926844\n",
      "\n",
      "results_dir: /Users/dave/.local/share/covalent/data\n",
      "dispatch_id: 0441d942-a8ab-4cb7-9373-1ef632ec395f\n",
      "\n",
      "Node Outputs\n",
      "------------\n",
      "task_1(0): 6\n",
      ":parameter:2(1): 2\n",
      ":sublattice:sub_workflow(2): [0, 1, 4, 9, 16, 25]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Technique 2: Correct\n",
    "\n",
    "# Define a sublattice that implements all the dynamic code\n",
    "@ct.electron\n",
    "@ct.lattice\n",
    "def sub_workflow(res):\n",
    "    \n",
    "    if res < 10:\n",
    "        final_res = []\n",
    "        for i in range(res):\n",
    "            final_res.append(task_2(i))\n",
    "    else:\n",
    "        final_res = res\n",
    "    \n",
    "    return final_res\n",
    "\n",
    "\n",
    "@ct.lattice\n",
    "def workflow_2(a):\n",
    "    res_1 = task_1(a)\n",
    "    return sub_workflow(res_1) # Nothing to see here. Just an electron consuming the output of another electron.\n",
    "\n",
    "id = ct.dispatch(workflow_2)(2)\n",
    "res = ct.get_result(id, wait=True)\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d105cc2b",
   "metadata": {},
   "source": [
    "## See Also\n",
    "\n",
    "[Result-Dependent Loops](./result_dependent_loop.ipynb)\n",
    "\n",
    "[Result-Dependent If-Else](./result_dependent_if_else.ipynb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
